# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: jhkim, taekwon

import copy
import collections
import recordclass

from coin.strategy.mm.tool.log_pta_dumper import OrderIdConverter

from coin.base.datetime_util import to_timestamp_int

from coin.proto.coin_order_gateway_pb2 import OrderGatewayLog, OrderEvent
import coin.proto.coin_order_gateway_pb2 as cogpb2
from coin.proto.coin_order_enums_pb2 import (BUY_ORDER,
                                             SELL_ORDER,
                                             BUY_OPEN_ORDER,
                                             BUY_CLOSE_ORDER,
                                             SELL_OPEN_ORDER,
                                             SELL_CLOSE_ORDER)

FillInfo = recordclass.recordclass('FillInfo',
                                   [
                                       'order_price',
                                       'sign',
                                       'order_qty',
                                       'fill_price',
                                       'fill_pq',
                                       'fill_qty',
                                       'submit_timestamp',
                                       'accepted_timestamp',
                                       'filled_timestamp',
                                   ])

OrderInfo = recordclass.recordclass('OrderInfo',
                                    [
                                        'order_id',
                                        'last_timestamp',
                                        'price',
                                        'sign',
                                        'order_qty',
                                        'fill_pq',
                                        'fill_qty',
                                        'cancelled',
                                        'submit_timestamp',
                                        'accepted_timestamp',
                                        'filled_timestamp',
                                        'cancel_submit_timestamp',
                                        'cancelled_timestamp',
                                        'rejected_timestamp',
                                        'rejected_perhaps_close'
                                    ])


def side_to_sign(side):
  if side in [SELL_ORDER, SELL_OPEN_ORDER, SELL_CLOSE_ORDER]:
    return -1
  elif side in [BUY_ORDER, BUY_OPEN_ORDER, BUY_CLOSE_ORDER]:
    return 1
  else:
    return 0


# per symbol


class OrderReader(object):
  def __init__(self, product):
    # order part
    self._orders_map = collections.OrderedDict()
    self._ext_int_map = {}
    self._bitflyer_submit_time = {}
    self._product = product
    self._order_id_converter = OrderIdConverter()
    # balance part
    self._balance_map = {}
    self._balances = collections.defaultdict(list)
    self._typescount = {sign: collections.defaultdict(int) for sign in [-1, 0, 1]}
    self._orders_list = []
    self._fill_list = []
    self._start_ts = 0
    self._end_ts = 0
    self._cleanup = False
    self._oes = []

  def set_time_range(self, start_ts, end_ts):
    self._start_ts = to_timestamp_int(start_ts or 0)
    self._end_ts = to_timestamp_int(end_ts or 0)

  def cleanup(self):
    if self._cleanup:
      return

    if self._product.exchange == 'Bitflyer':
      for external_id in self._orders_map.keys():
        internal_id = self._ext_int_map.get(external_id, None)
        if internal_id in self._bitflyer_submit_time:
          self._orders_map[external_id].submit_timestamp = self._bitflyer_submit_time[internal_id]
      for each_order in self._orders_list:
        if each_order.order_id in self._orders_map:
          each_order.submit_timestamp = self._orders_map[each_order.order_id].submit_timestamp

  def get_orders_map(self):
    self.cleanup()
    return self._orders_map

  def get_orders_list(self):
    self.cleanup()
    return self._orders_list

  def get_balances(self, key):
    print(key)
    print(self._balances.keys())
    ts, bals = zip(*self._balances[key])
    tsbals_chosen = []
    for ts_iter, bal_iter in zip(ts, bals):
      if len(tsbals_chosen) > 0 and tsbals_chosen[-1][0] > ts_iter:
        continue
      tsbals_chosen.append((ts_iter, bal_iter))
    ts, bals = zip(*tsbals_chosen)
    return ts, bals

  def on_log2(self, timestamp, log):
    if self._start_ts and timestamp < self._start_ts:
      return
    if self._end_ts and timestamp > self._end_ts:
      return
    pb = OrderGatewayLog()
    try:
      pb.ParseFromString(log)
    except Exception:
      print(log)
      return

    if pb.type == OrderGatewayLog.ORDER_EVENT:
      self.on_orderevent_pb(timestamp, pb)
    elif pb.type == OrderGatewayLog.BALANCE:
      self.on_balance_pb(timestamp, pb)
    elif pb.type == OrderGatewayLog.POSITION:
      self.on_position_pb(timestamp, pb)

  def on_position_pb(self, timestamp, pb):
    if pb.type != OrderGatewayLog.POSITION:
      return
    position = pb.position
    for each_position in position.each_position:
      symbol = each_position.symbol
      self._balance_map[symbol] = each_position.net_position
      self._balances[symbol].append((timestamp, self._balance_map[symbol]))

  def on_balance_pb(self, timestamp, pb):
    if pb.type != OrderGatewayLog.BALANCE:
      return
    balance = pb.balance
    for each_balance in balance.each_balance:
      self._balance_map[each_balance.currency] = each_balance.total
      self._balances[each_balance.currency].append(
          (timestamp, self._balance_map[each_balance.currency]))

  def on_orderevent_pb(self, timestamp, pb):
    oe = pb.event
    if not oe.HasField('internal_order_id'):
      oe.internal_order_id = str(oe.proc_order_id)
    etype_str = cogpb2._ORDEREVENT_ORDEREVENTTYPE.values_by_number[oe.type].name
    if pb.type != OrderGatewayLog.ORDER_EVENT:
      return
    sign = side_to_sign(oe.order_side)
    self._typescount[sign][etype_str] += 1

    if oe.type not in [
        OrderEvent.ORDER_SUBMITTED,
        OrderEvent.ORDER_AMENDED,
        OrderEvent.ORDER_AMEND_ACCEPTED,
        OrderEvent.ORDER_ACCEPTED,
        OrderEvent.ORDER_AUTO_CANCELED,
        OrderEvent.ORDER_FILLED,
        OrderEvent.ORDER_ERROR,
        OrderEvent.ORDER_REJECTED,
        OrderEvent.CANCEL_CONFIRMED,
        OrderEvent.CANCEL_SUBMITTED,
        OrderEvent.CANCEL_ACCEPTED]:
      return

    if oe.symbol != self._product.symbol:
      return

    self._oes.append(oe)

    if oe.type in (OrderEvent.ORDER_REJECTED, OrderEvent.ORDER_ERROR):
      order_id = self._order_id_converter.get_order_id_from_pb(oe)

      if order_id in self._orders_map:
        self._orders_map[order_id].rejected_timestamp = timestamp
        self._typescount[self._orders_map[order_id].sign][etype_str] += 1
        if oe.order_side in (BUY_CLOSE_ORDER, SELL_CLOSE_ORDER):
          self._orders_map[order_id].rejected_perhaps_close = True
        else:
          self._orders_map[order_id].rejected_perhaps_close = False
        self._orders_list.append(copy.deepcopy(self._orders_map[order_id]))

    if self._product.exchange == 'Bitflyer':
      tshash = int(timestamp / 1e11)
      if oe.HasField('external_order_id'):
        order_id = oe.external_order_id
        if oe.HasField('internal_order_id'):
          self._ext_int_map[oe.external_order_id] = (tshash, oe.internal_order_id)
      else:
        if oe.type == OrderEvent.ORDER_SUBMITTED:
          self._bitflyer_submit_time[(tshash, oe.internal_order_id)] = timestamp
        return
    else:
      order_id = self._order_id_converter.get_order_id_from_pb(oe)
      if oe.type == OrderEvent.ORDER_AMEND_ACCEPTED:
        if order_id in self._orders_map:
          self._orders_list.append(copy.deepcopy(self._orders_map[order_id]))
          self._orders_list[-1].cancelled = True
        else:
          order_id = 0
        self._orders_list.append(
            OrderInfo(order_id or 0 + 1e8,
                      -1,
                      None,
                      None,
                      0,
                      0,
                      0,
                      False,
                      -1,
                      -1,
                      -1,
                      -1,
                      -1,
                      -1,
                      None))
        self._orders_list[-1].price = oe.order_price
        self._orders_list[-1].order_qty = oe.order_qty
        self._orders_list[-1].sign = sign
        self._orders_list[-1].last_timestamp = timestamp
        self._orders_list[-1].submit_timestamp = timestamp
      if order_id not in self._orders_map:
        if oe.type not in (OrderEvent.ORDER_SUBMITTED, OrderEvent.ORDER_AMEND_ACCEPTED):
          return

    if ((self._product.exchange == 'Bitflyer' and order_id not in self._orders_map)
        or (oe.type == OrderEvent.ORDER_SUBMITTED) or (oe.type == OrderEvent.ORDER_AMEND_ACCEPTED)):
      self._orders_map[order_id] = OrderInfo(order_id,
                                             -1,
                                             None,
                                             None,
                                             0,
                                             0,
                                             0,
                                             False,
                                             -1,
                                             -1,
                                             -1,
                                             -1,
                                             -1,
                                             -1,
                                             None)

    self._orders_map[order_id].last_timestamp = timestamp

    if oe.type == OrderEvent.CANCEL_SUBMITTED:
      self._orders_map[order_id].cancel_submit_timestamp = timestamp

    if oe.type == OrderEvent.CANCEL_CONFIRMED:
      self._orders_map[order_id].cancelled = True
      self._orders_map[order_id].cancelled_timestamp = timestamp
      self._orders_list.append(copy.deepcopy(self._orders_map[order_id]))
      self._orders_list[-1].fill_pq = 0
      self._orders_list[-1].fill_qty = 0

    if oe.type == OrderEvent.ORDER_SUBMITTED:
      self._orders_map[order_id].submit_timestamp = timestamp

    if (oe.type == OrderEvent.ORDER_SUBMITTED or (oe.type == OrderEvent.ORDER_AMEND_ACCEPTED)
        or (oe.type == OrderEvent.ORDER_ACCEPTED and self._product.exchange == 'Bitflyer')):
      self._orders_map[order_id].price = oe.order_price
      self._orders_map[order_id].order_qty = oe.order_qty
      self._orders_map[order_id].sign = sign
      self._orders_list.append(copy.deepcopy(self._orders_map[order_id]))
      self._orders_list[-1].fill_pq = 0
      self._orders_list[-1].fill_qty = 0

    if oe.type == OrderEvent.ORDER_ACCEPTED or (oe.type == OrderEvent.ORDER_AMEND_ACCEPTED):
      self._orders_map[order_id].accepted_timestamp = timestamp
      # overload sign. cuz submit sign is not trustworthy
      self._orders_map[order_id].sign = sign
    elif oe.type == OrderEvent.ORDER_FILLED:
      self._orders_map[order_id].filled_timestamp = timestamp
      self._orders_map[order_id].fill_pq += oe.fill_price * oe.fill_qty
      self._orders_map[order_id].fill_qty += oe.fill_qty

      self._orders_list.append(copy.deepcopy(self._orders_map[order_id]))
      # duplicated fills!
      self._orders_list[-1].fill_pq = oe.fill_price * oe.fill_qty
      self._orders_list[-1].fill_qty = oe.fill_qty

      # 'price',
      # 'sign',
      # 'order_qty',
      # 'fill_pq',
      # 'fill_qty',
      # 'submit_timestamp',
      # 'accepted_timestamp',
      # 'filled_timestamp',
      self._fill_list.append(
          FillInfo(self._orders_map[order_id].price,
                   self._orders_map[order_id].sign,
                   self._orders_map[order_id].order_qty,
                   oe.fill_price,
                   oe.fill_price * oe.fill_qty,
                   oe.fill_qty,
                   self._orders_map[order_id].submit_timestamp,
                   self._orders_map[order_id].accepted_timestamp,
                   timestamp))
